from PIL import Image

def arnold_encrypt_image(image_path, rounds):
    # 打开图像
    image = Image.open(image_path).convert("L")  # 转换为灰度图像
    width, height = image.size

    # 迭代加密轮数
    for _ in range(rounds):
        # 保存临时结果
        temp_image = Image.new("L", (width, height))
        # 替换和置换像素
        for y in range(height):
            for x in range(width):
                new_x = (x + y) % width
                new_y = (x + 2 * y) % height
                temp_image.putpixel((new_x, new_y), image.getpixel((x, y)))
        image = temp_image

    return image

def arnold_decrypt_image(image, rounds):
    width, height = image.size
    # 初始化解密后的图像
    decrypted_image = Image.new("L", (width, height))

    # 迭代解密轮数
    for _ in range(rounds):
        # 保存临时结果
        temp_image = Image.new("L", (width, height))
        # 替换和置换像素
        for y in range(height):
            for x in range(width):
                new_x = (2 * x - y) % width
                new_y = (-x + y) % height
                temp_image.putpixel((new_x, new_y), image.getpixel((x, y)))
        image = temp_image

    return image

# 测试
input_image_path = "Lena.bmp"
output_image_path = "encrypted_image_gray4.jpg"
decrypted_image_path = "decrypted_image_gray4.jpg"
rounds = 10

# 加密图像
encrypted_image = arnold_encrypt_image(input_image_path, rounds)
encrypted_image.save(output_image_path)

# 解密图像
decrypted_image = arnold_decrypt_image(encrypted_image, rounds)
decrypted_image.save(decrypted_image_path)